import hashlib
import json
import os
import sys
import time
from pathlib import Path

from PyQt5.QtCore import QTimer
from PyQt5.QtWidgets import QApplication

from launch import launch_windows_app
from upgrade import UpgradeUI, AsyncWorker


# pyinstaller -F main.py -n upgrade;pyinstaller .\upgrade.spec
# this main block is required to generate executable by pyinstaller

def calculate_md5(file_path, chunk_size=8192):
    """计算文件的MD5哈希值"""
    hash_obj = hashlib.md5()

    with open(file_path, 'rb') as f:
        while chunk := f.read(chunk_size):
            hash_obj.update(chunk)  # 更新哈希值

    return hash_obj.hexdigest()  # 返回十六进制字符串

def check_model_integrity(filepath, checksum):
    if os.path.exists(filepath):
        return calculate_md5(filepath) == checksum
    return False

def upgrade_image(upg: UpgradeUI, localdir, upgradelist):
    # upgradelist = ['u2net/u2net-latest.txt', 'sam-hq/sam-hq-latest.txt', 'automatting/automatting-lastest.txt']
    # upgradelist = ['u2net/u2net-latest.txt']
    try:
        for desc in upgradelist:
            image, imagelatest = desc.split('/')
            if localdir == '':
                local_dir = ''
                local_imagelatest = f'{imagelatest}'
            else:
                local_dir = f'{localdir}/{image}'
                local_imagelatest = f'{localdir}/{image}/{imagelatest}'

            Path(local_dir).mkdir(parents=True, exist_ok=True)

            imageconfigs = {
                'url':[],
                'checksum':[],
                'filename':[],
                'image':[],
            }
            try:
                # check local model updates
                with open(local_imagelatest, "r") as f:
                    imageconfigs = json.load(f)
            except Exception as e:
                pass

            # worker = AsyncWorker(upg.download_description, image, imagelatest)
            # worker.start()
            # worker.wait()
            # if not worker.result:
            #     raise Exception('download error')
            if upg.download_description(image, imagelatest) is False:
                raise Exception('download error')

            print(upg.imageconfigs)
            needupdates = False

            try:
                if len(imageconfigs['filename']) != 0:
                    for i, filename in enumerate(imageconfigs['filename']):
                        if localdir == '':
                            filepath = f'{filename}'
                        else:
                            filepath = f'{local_dir}/{filename}'
                        checksum = imageconfigs['checksum'][i]
                        if check_model_integrity(filepath, checksum) is False:
                            needupdates = True
                            break
                else:
                    needupdates = True
            except Exception as e:
                needupdates = True

            for checksum in imageconfigs['checksum']:
                if checksum not in upg.imageconfigs['checksum']:
                    # Need updates
                    needupdates = True

            if needupdates:
                for i, filename in enumerate(upg.imageconfigs['filename']):
                    url = upg.imageconfigs['url'][i]
                    checksum = upg.imageconfigs['checksum'][i]
                    # worker = AsyncWorker(upg.download_file, url, local_dir)
                    # worker.start()
                    # worker.wait()
                    # if not worker.result:
                    #     raise Exception("download error")
                    if upg.download_file(url, local_dir) is False:
                        raise Exception("download error")
                    filepath = os.path.join(local_dir, os.path.basename(filename))
                    check_model_integrity(filepath, checksum)
                try:
                    # updates modellatest file locally
                    with open(local_imagelatest, 'w') as f:
                        json.dump(upg.imageconfigs, f, indent=2)
                except Exception as e:
                    pass
    except Exception as e:
        return False

    return True

def upgrade(upg):
    print("upgrade ...")
    upgradelist = ['u2net/u2net-latest.txt', 'sam-hq/sam-hq-latest.txt']
    upgrade_image(upg, 'models', upgradelist)
    # worker = AsyncWorker(upgrade_image, upg, 'models', upgradelist)
    # worker.start()
    # worker.wait()
    # print(f'worker.result = {worker.result}')
    #
    upgradelist = ['automatting/automatting-latest.txt']
    upgrade_image(upg, '', upgradelist)
    # worker = AsyncWorker(upgrade_image, upg, '', upgradelist)
    # worker.finished.connect(app.quit)
    # worker.start()
    # worker.wait()
    # upg.dialog_close.emit()
    print("launching apps")
    upg.label_updated.emit('Launching applications ...')
    launch_windows_app(upg, './automatting.exe', ['--no-upgrade'], False, './')

def redirect_output():
    if hasattr(sys, 'frozen') and not hasattr(sys, 'ps1'):
        # 检查是否是打包环境且无控制台
        output_dir = os.path.dirname(sys.executable)
        stdout_path = os.path.join(output_dir, 'stdout.log')
        stderr_path = os.path.join(output_dir, 'stderr.log')

        sys.stdout = open(stdout_path, 'w', encoding='utf-8')
        sys.stderr = open(stderr_path, 'w', encoding='utf-8')

if __name__ == '__main__':
    redirect_output()
    app = QApplication(sys.argv)
    upg = UpgradeUI()
    upg.show()
    worker = AsyncWorker(upgrade, upg)
    worker.start()

    sys.exit(app.exec_())